#version 450 core
#extension GL_ARB_explicit_uniform_location : require
#define NUM_MAX_DIRLIGHTS 8
#define NUM_MAX_POINTLIGHTS 8
#define NUM_MAX_SPOTLIGHTS 8

layout(location = 0) in vec3 normal;
layout(location = 1) in vec3 fragPos;
layout(location = 0) out vec4 fragColor;

struct World {
    vec3 cameraPos;
    mat4 matVP;
};
struct Transform {
    mat4 modelMatrix;
    mat4 mvpMatrix;
    mat4 normalMatrix;
};
struct Material {
    vec3 color;
    float shiness;
};
struct DirLight {
    vec3 ambient;
    vec3 diffuse;
    vec3 specular;
    vec3 direction;
};
struct PointLight {
    vec3 ambient;
    vec3 diffuse;
    vec3 specular;
    vec3 position;
    float c0, c1, c2;
};
struct SpotLight {
    vec3 color;
    vec3 position;
    vec4 data0; // data0.xyz: direction, data0.w: cutOff (in degrees)
    vec4 data1; // data1.xyz: ambient, data1.w: outerCutOff
};
layout(binding = 0, std140) uniform WorldBlock {
    World world;
};
layout(binding = 1, std140) uniform DirLightBlock {
    DirLight dirLight[NUM_MAX_DIRLIGHTS];
};
layout(binding = 2, std140) uniform PointLightBlock {
    PointLight pointLight[NUM_MAX_POINTLIGHTS];
};
layout(binding = 3, std140) uniform SpotLightBlock {
    SpotLight spotLight[NUM_MAX_SPOTLIGHTS];
};
layout(location = 0) uniform Transform transform;
layout(location = 3) uniform Material material;
layout(location = 5) uniform uint numDirLights;
layout(location = 6) uniform uint numPointLights;
layout(location = 7) uniform uint numSpotLights;

void main() {
    vec3 result = vec3(0.0);
    const vec3 viewDir = normalize(world.cameraPos - fragPos);
    for (uint i = 0; i < numDirLights; ++i) {
        vec3 ambient = dirLight[i].ambient * material.color;
        vec3 lightDir = -normalize(dirLight[i].direction);
        vec3 diffuse = max(dot(lightDir, normal), 0.0f) * dirLight[i].diffuse;
        vec3 reflectDir = reflect(-lightDir, normal);
        float spec = pow(max(dot(viewDir, reflectDir), 0.0f), material.shiness);
        vec3 specular = spec * dirLight[i].specular * material.color;
        result += (ambient + diffuse + specular);
    }
    for (uint i = 0; i < numPointLights; ++i) {
        float dis = distance(pointLight[i].position, fragPos);
        float attenuation = pointLight[i].c0 + pointLight[i].c1 * dis + pointLight[i].c2 * dis * dis;
        vec3 ambient = pointLight[i].ambient * material.color;
        vec3 lightDir = normalize(pointLight[i].position - fragPos);
        vec3 diffuse = max(dot(lightDir, normal), 0.0f) * pointLight[i].diffuse;
        vec3 reflectDir = reflect(-lightDir, normal);
        float spec = pow(max(dot(viewDir, reflectDir), 0.0f), material.shiness);
        vec3 specular = spec * pointLight[i].specular * material.color;
        result += (ambient + diffuse + specular) / attenuation;
    }
    for (uint i = 0; i < numSpotLights; ++i) {
        vec3 lightDir = normalize(spotLight[i].position - fragPos);
        float c0 = cos(radians(spotLight[i].data0.w)); // inner
        float c1 = cos(radians(spotLight[i].data1.w)); // outer
        float cosTheta = dot(lightDir, -normalize(spotLight[i].data0.xyz));
        float intensity = clamp((cosTheta - c1) / (c0 - c1), 0.0f, 1.0f);
        vec3 diffuse = max(dot(lightDir, normal), 0.0f) * spotLight[i].color;
        vec3 reflectDir = reflect(-lightDir, normal);
        float spec = pow(max(dot(viewDir, reflectDir), 0.0f), material.shiness);
        vec3 specular = spec * spotLight[i].color * material.color;
        result += (spotLight[i].data1.xyz + (diffuse + specular) * intensity); // * attenuation
    }
    fragColor = vec4(result, 1.0f);
}
